
from jieba import strdecode, resolve_filename
from jieba.finalseg import emit_P
from jieba.posseg import re_han_internal, re_skip_internal, pair, re_num, re_eng, char_state_tab_P, start_P, trans_P
from jieba.posseg.viterbi import MIN_FLOAT, MIN_INF, viterbi


class PosTagging(object):
    def __cut_internal(self, sentence, HMM=True):
        # 如果没有字典没有初始化，则先加载字典。否则直接使用字典缓存即可。
        self.makesure_userdict_loaded()

        # 将语句转为UTF-8或者GBK
        sentence = strdecode(sentence)

        # 根据正则匹配，将输入文本分隔成一个个语句。分隔符包括空格 逗号 句号等。
        blocks = re_han_internal.split(sentence)

        # 根据是否采用了HMM模型来进行不同方法的选择
        if HMM:
            cut_blk = self.__cut_DAG
        else:
            cut_blk = self.__cut_DAG_NO_HMM

        # 遍历正则匹配分隔好的语句，对每个语句进行单独的分词和词性标注
        for blk in blocks:
            if re_han_internal.match(blk):
                # 分词和词性标注
                for word in cut_blk(blk):
                    yield word
            else:
                tmp = re_skip_internal.split(blk)
                for x in tmp:
                    if re_skip_internal.match(x):
                        yield pair(x, 'x')
                    else:
                        for xx in x:
                            if re_num.match(xx):
                                yield pair(xx, 'm')
                            elif re_eng.match(x):
                                yield pair(xx, 'eng')
                            else:
                                yield pair(xx, 'x')

    def __cut_DAG(self, sentence):
        # 构建DAG有向无环图，得到语句分词所有可能的路径
        DAG = self.tokenizer.get_DAG(sentence)
        route = {}
        # 动态规划，计算从语句末尾到语句起始，DAG中每个节点到语句结束位置的最大路径概率，以及概率最大时节点对应词语的结束位置
        self.tokenizer.calc(sentence, DAG, route)
        # 遍历每个节点的Route，组装词语组合。
        x = 0
        buf = ''
        N = len(sentence)
        while x < N:
            # y表示词语的结束位置，x为词语的起始位置
            y = route[x][1] + 1
            # 从起始位置x到结束位置y，取出一个词语
            l_word = sentence[x:y]
            if y - x == 1:
                # 单字，一个汉字构成的一个词语
                buf += l_word
            else:
                # 多汉字词语
                if buf:
                    if len(buf) == 1:
                        # 单字直接从字典中取出其词性。使用pair将分词和词性一起输出。
                        yield pair(buf, self.word_tag_tab.get(buf, 'x'))
                    elif not self.tokenizer.FREQ.get(buf):
                        # 词语不在字典中，也就是新词，使用HMM隐马尔科夫模型进行分割
                        recognized = self.__cut_detail(buf)
                        for t in recognized:
                            yield t
                    else:
                        # 词语在字典中，直接查找字典并取出词性。
                        for elem in buf:
                            yield pair(elem, self.word_tag_tab.get(elem, 'x'))
                    buf = ''
                yield pair(l_word, self.word_tag_tab.get(l_word, 'x'))
            # 该节点取词完毕，跳到下一个词语的开始位置
            x = y
        # 通过yield，逐词返回上一步切分好的词语
        if buf:
            if len(buf) == 1:
                yield pair(buf, self.word_tag_tab.get(buf, 'x'))
            elif not self.tokenizer.FREQ.get(buf):
                recognized = self.__cut_detail(buf)
                for t in recognized:
                    yield t
            else:
                for elem in buf:
                    yield pair(elem, self.word_tag_tab.get(elem, 'x'))

        def load_word_tag(self, f):
            self.word_tag_tab = {}
            f_name = resolve_filename(f)
            # 遍历字典的每一行。每一行对应一个词语。包含词语 词数 词性三部分
            for lineno, line in enumerate(f, 1):
                try:
                    # 去除首尾空格符
                    line = line.strip().decode("utf-8")
                    if not line:
                        continue
                    # 利用空格将一行分隔为词语 词数 词性三部分
                    word, _, tag = line.split(" ")
                    # 使用词语为key，词性为value，构造Dict
                    self.word_tag_tab[word] = tag
                except Exception:
                    raise ValueError(
                        'invalid POS dictionary entry in %s at Line %s: %s' % (f_name, lineno, line))
            f.close()

    # 通过HMM隐马尔科夫模型获取词性标注序列，解决未登录的问题
    def __cut(self, sentence):
        # 通过viterbi算法，利用三个概率，由语句观测序列，得到词性标注隐藏序列
        # prob为
        # pos_list对应每个汉字，包含分词标注BEMS和词语词性两部分。
        prob, pos_list = viterbi(sentence, char_state_tab_P, start_P, trans_P, emit_P)
        begin, nexti = 0, 0
        # 遍历语句的每个汉字，如果是E或者S时，也就是词语结束或者单字词语，则分隔得到词语和词性pair
        for i, char in enumerate(sentence):
            pos = pos_list[i][0]
            if pos == 'B':
                # B表示词语的开始
                begin = i
            elif pos == 'E':
                # E表示词语的结束，此时输出词语和他的词性
                yield pair(sentence[begin:i + 1], pos_list[i][1])
                nexti = i + 1
            elif pos == 'S':
                # S表示单字词语，此时也输出词语和他的词性
                yield pair(char, pos_list[i][1])
                nexti = i + 1
        # 一般不会走到这儿，以防万一。对剩余的所有汉字一起输出一个词语和词性。
        if nexti < len(sentence):
            yield pair(sentence[nexti:], pos_list[nexti][1])

    def viterbi(obs, states, start_p, trans_p, emit_p):
        V = [{}]  # 每个汉字的每个BEMS状态的最大概率。
        mem_path = [{}]  # 分词路径
        # 初始化每个state，states为"BEMS"
        all_states = trans_p.keys()
        for y in states.get(obs[0], all_states):  # init
            V[0][y] = start_p[y] + emit_p[y].get(obs[0], MIN_FLOAT)
            mem_path[0][y] = ''
        # 逐字进行处理
        for t in range(1, len(obs)):
            V.append({})
            mem_path.append({})
            # prev_states = get_top_states(V[t-1])
            prev_states = [
                x for x in mem_path[t - 1].keys() if len(trans_p[x]) > 0]
            prev_states_expect_next = set(
                (y for x in prev_states for y in trans_p[x].keys()))
            obs_states = set(
                states.get(obs[t], all_states)) & prev_states_expect_next
            if not obs_states:
                obs_states = prev_states_expect_next if prev_states_expect_next else all_states
            # 遍历每个状态
            for y in obs_states:
                # 计算前一个状态到本状态的最大概率和它的前一个状态
                prob, state = max((V[t - 1][y0] + trans_p[y0].get(y, MIN_INF) +
                                   emit_p[y].get(obs[t], MIN_FLOAT), y0) for y0 in prev_states)
                # 将该汉字下的某状态（BEMS）的最大概率记下来
                V[t][y] = prob
                # 记录状态转换路径
                mem_path[t][y] = state
        last = [(V[-1][y], y) for y in mem_path[-1].keys()]
        # if len(last)==0:
        #     print obs
        prob, state = max(last)
        route = [None] * len(obs)
        i = len(obs) - 1
        while i > 0:
            route[i] = state
            state = mem_path[i][state]
            i -= 1
        return (prob, route)

